"""
Visualization utilities for LLM concepts and results.
"""

import matplotlib.pyplot as plt
import numpy as np


def plot_attention_weights(attention_weights, tokens=None, title="Attention Weights"):
    """
    Plot attention weights as a heatmap.
    
    Args:
        attention_weights: 2D array of attention weights
        tokens: Optional list of token labels
        title: Plot title
    """
    fig, ax = plt.subplots(figsize=(10, 8))
    im = ax.imshow(attention_weights, cmap='viridis')
    
    if tokens:
        ax.set_xticks(np.arange(len(tokens)))
        ax.set_yticks(np.arange(len(tokens)))
        ax.set_xticklabels(tokens)
        ax.set_yticklabels(tokens)
        plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
    
    ax.set_title(title)
    fig.colorbar(im, ax=ax)
    plt.tight_layout()
    return fig, ax


def plot_training_curves(train_losses, val_losses=None, title="Training Curves"):
    """
    Plot training and validation loss curves.
    
    Args:
        train_losses: List of training losses
        val_losses: Optional list of validation losses
        title: Plot title
    """
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(train_losses, label='Training Loss')
    
    if val_losses:
        ax.plot(val_losses, label='Validation Loss')
    
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss')
    ax.set_title(title)
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    return fig, ax

